"港航 AI 识别完毕后处理代码"

__author__ = "zwj"

# 导入根目录作为模块搜索路径
import os, sys

# ([".."] * 4) 就是项目目录里面 4 层
parent_path = os.path.abspath(os.path.join(__file__, *([".."] * 4)))
sys.path.insert(0, parent_path)

import yaml, time, cv2
from typing import Dict, List
from enum import Enum, unique

from core.task.modules.processors.base_processor import BaseProcessor
from core.utils.logger import logging


# 船只航向枚举
@unique
class Heading(Enum):
    kLeft = "left"
    kRight = "right"
    kUnknown = "unknown"


# 像素点坐标
class PointData:
    def __init__(self, x=0, y=0):
        # 确保点坐标都为正数
        self.x = abs(x)
        self.y = abs(y)


# 船只检测数据汇总
class ShipData:
    def __init__(
        self,
        cls_id,
        track_id,
        ship_start: PointData,
        ship_end: PointData,
    ):
        # 船只检测
        self.cls_id = cls_id
        self.track_id = track_id
        self.ship_start = ship_start
        self.ship_end = ship_end

    def get_midpoint_data(self) -> PointData:
        return PointData(
            (self.ship_start.x + self.ship_end.x) / 2,
            (self.ship_start.y + self.ship_end.y) / 2,
        )

    def get_area(self) -> int:
        return (self.ship_end.x - self.ship_start.x) * (
            self.ship_end.y - self.ship_start.y
        )


# 船只检测数据
class Mask:
    def __init__(self, width, height):
        self.__width = width
        self.__height = height
        # 初始化一个二维列表
        self.clear()

    def set_mask(self, start_point: PointData, end_point: PointData):
        # 确保数据不过界
        start_point = self.__point_data_int(start_point)
        end_point = self.__point_data_int(end_point)

        # 将给定框范围内的元素置为True
        for y in range(start_point.y, end_point.y):
            for x in range(start_point.x, end_point.x):
                self.__mask[y][x] = True

    def get_mask(self, point: PointData):
        point = self.__point_data_int(point)
        return self.__mask[int(point.y)][int(point.x)]

    def clear(self):
        self.__mask = [
            [False for _ in range(self.__width)] for _ in range(self.__height)
        ]

    def __point_data_int(self, point: PointData):
        point_x = min(int(point.x), self.__width - 1)
        point_y = min(int(point.y), self.__height - 1)
        return PointData(int(point_x), int(point_y))


@unique
class StateMachineEnum(Enum):
    kDetect = 0  # 检测船只,锁定要处理的船只
    kZoom = 1  # 变焦
    kZoomDetect = 2  # 变焦后检测


class GangHangProcessor(BaseProcessor):
    def __init__(
        self,
        keys,
        Custom_cls_map,
        Screen_size,
        Left_Heading,
    ) -> None:

        # * key 值对应索引
        self.__keys = keys

        # * 输入数据
        self.__in_key = self.__keys["in"]
        self.__track_key = self.__keys["track"]
        self.__out_key = self.__keys["out"]
        self.__finish_tag_key = self.__keys["finish_tag"]

        # * 检测 id 对应含义地图
        if isinstance(Custom_cls_map, str):
            with open(Custom_cls_map, "r") as f:
                self.__cls_map = {
                    x: y for x, y in enumerate(yaml.safe_load(f)["label_list"])
                }
        else:
            self.__cls_map = Custom_cls_map

        # * 绘图框颜色初始化, BGR 颜色, 0 0 0 为黑色
        self.__color = (0, 0, 255)

        # * 图片输出目录初始化
        self.__image_output_path = os.getcwd() + "/output"
        if not os.path.exists(self.__image_output_path):
            os.makedirs(self.__image_output_path)

        # * 获取画面大小
        self.__screen_size_x = Screen_size[0]
        self.__screen_size_y = Screen_size[1]

        # * 船只数据
        self.__ship_data_list: List[ShipData] = []
        self.__detecting_ships_list: List[ShipData] = []

        # * 船只变焦下截图
        self.__zoom_image_flag = False
        self.__zoom_image_time = time.time()

        # 船只掩码
        self.__mask = Mask(self.__screen_size_x, self.__screen_size_y)

        # 船只航向要求
        self.__set_ship_heading(Left_Heading)

        # * 记录操作时间戳
        self.__last_fps_time = time.time()

    def __set_ship_heading(self, Left_Heading):
        # 初始化一个掩码二维列表,将不检测区屏蔽
        self.__mask.clear()
        self.__mask.set_mask(
            PointData(0, 0),
            PointData(self.__screen_size_x, self.__screen_size_y // 3),
        )

        if Left_Heading:
            self.__ship_heading = Heading.kLeft
            self.__cls_id = 2

            self.__mask.set_mask(
                PointData(self.__screen_size_x // 2, 0),
                PointData(self.__screen_size_x, self.__screen_size_y),
            )
        else:
            self.__ship_heading = Heading.kRight
            self.__cls_id = 0

            self.__mask.set_mask(
                PointData(0, 0),
                PointData(self.__screen_size_x // 2, self.__screen_size_y),
            )

    def __refresh_ship_data(self, data: Dict) -> Dict:
        # 获取结果
        track_result = data[self.__track_key]

        ship_data_list = []
        detecting_ships_list = []

        for bboxes, track_ids in zip(track_result["boxes"], track_result["track_ids"]):

            for box, track_id in zip(bboxes, track_ids):
                cls_id, score, x1, y1, x2, y2 = box
                cls_id = int(cls_id)

                # 只记录检测区域内船只
                ship_data = ShipData(
                    cls_id, track_id, PointData(x1, y1), PointData(x2, y2)
                )

                ship_data_list.append(ship_data)

                if cls_id != self.__cls_id:
                    continue

                # 获取到的船中心点如果在掩码区域内跳过
                if self.__mask.get_mask(ship_data.get_midpoint_data()):
                    continue

                detecting_ships_list.append(ship_data)

        self.__ship_data_list = ship_data_list
        self.__detecting_ships_list = detecting_ships_list

        return data

    def __detect_image(self, data: Dict, ships_list: List[ShipData]) -> Dict:
        image = data[self.__in_key][0].copy()
        draw_thickness = min(image.shape[:2]) // 320
        arrow_length = 40  # 箭头的长度

        if self.__zoom_image_flag and time.time() - self.__zoom_image_time > 1:
            self.__zoom_image_time = time.time()
            cv2.imwrite(f"{self.__image_output_path}/{time.time()}.jpg", image)

        for ship_data in ships_list:
            cls_id = ship_data.cls_id

            if ship_data.cls_id > 4:
                color = (255, 255, 255)
                cls_id -= 5
            else:
                color = self.__color

            midpoint_data = ship_data.get_midpoint_data()

            arrow_start = (
                int(midpoint_data.x),
                int(midpoint_data.y),
            )  # 箭头的起点坐标

            # 画出检测类型
            cv2.putText(
                image,
                self.__cls_map[ship_data.cls_id],
                (int(ship_data.ship_start.x), int(ship_data.ship_start.y)),
                cv2.FONT_HERSHEY_COMPLEX,
                1,
                color,
                2,
            )
            # 画出追踪框
            cv2.rectangle(
                image,
                (int(ship_data.ship_start.x), int(ship_data.ship_start.y)),
                (int(ship_data.ship_end.x), int(ship_data.ship_end.y)),
                color=color,
                thickness=draw_thickness,
            )

            # 画出检测类型
            if cls_id < 4:
                if cls_id == 0:
                    # 右箭头
                    arrow_end = (arrow_start[0] + arrow_length, arrow_start[1])
                elif cls_id == 1:
                    # 下箭头
                    arrow_end = (arrow_start[0], arrow_start[1] + arrow_length)
                elif cls_id == 2:
                    # 左箭头
                    arrow_end = (arrow_start[0] - arrow_length, arrow_start[1])
                elif cls_id == 3:
                    # 上箭头
                    arrow_end = (arrow_start[0], arrow_start[1] - arrow_length)

                # 白底黑箭头
                cv2.arrowedLine(
                    image,
                    arrow_start,
                    arrow_end,
                    (255, 255, 255),
                    draw_thickness,
                )
                cv2.arrowedLine(
                    image,
                    arrow_start,
                    arrow_end,
                    (0, 0, 0),
                    draw_thickness // 2,
                )

        self.__out.write(image)

        data[self.__out_key] = image
        return data

    def _call(self, data: Dict) -> Dict:
        # 计算出调用该函数的帧率
        # fps = 1 / (time.time() - self.__last_fps_time)
        # logging.info(f"帧率 : {fps}")
        # self.__last_fps_time = time.time()

        # 获取船只数据
        data = self.__refresh_ship_data(data)

        # 画出状态机流程的检测框
        # data = self.__detect_image(data, self.__detecting_ships_list)

        # 画出所有检测船只
        data = self.__detect_image(data, self.__ship_data_list)

        # 读取完毕关闭任务
        if data.get(self.__finish_tag_key, False):
            logging.info("任务完成")
            self.close()

        return data

    def _init_check(self):
        """
        任务执行前初始化检查参数
        """

        mkv_path = os.path.join(self.__image_output_path, "output.mkv")
        self.__out = cv2.VideoWriter(
            mkv_path,
            cv2.VideoWriter_fourcc(*"XVID"),
            30,
            (self.__screen_size_x, self.__screen_size_y),
        )

    def _set(self, data: Dict):
        if data.get("Left_Heading", False):
            self.__set_ship_heading(data["Left_Heading"])

    def _start(self):
        pass

    def _close(self):
        self.__out.release()


if __name__ == "__main__":
    pass
